#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient, IndexModel, ASCENDING, DESCENDING
import re


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


# 第一步导入数据后清除空字段
def Clean_eLMSG():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")

    for result in collection_eLMSG.find():
        del_empty_feild(result)
        collection_eLMSG.replace_one({"_id": result["_id"]}, result)

    CloseMongoDB(client_eLMSG)


def del_empty_feild(_dict: dict, parent=None, key_in_parent=None):
    for k in list(_dict.keys()):
        v = _dict[k]
        if isinstance(v, str):
            if v is None or not len(v.strip()):
                del _dict[k]
        elif isinstance(v, list):
            if not len(v):
                del _dict[k]
            else:
                for i in v:
                    if isinstance(i, dict):
                        del_empty_feild(i)
        elif isinstance(v, dict):
            del_empty_feild(v, _dict, k)
        elif v is None:
            del _dict[k]

    if not len(_dict) and parent is not None and key_in_parent is not None:
        del parent[key_in_parent]


# 生成结构化的Reference
def GenerateStructuredReference():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")

    for result in collection_eLMSG.find({"References": {"$exists": 1}}, {"References": 1}):
        reference = []
        for line in result["References"]:
            reference.append(GetReference(line))
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"StructuredReference": reference}})
    CloseMongoDB(client_eLMSG)


def GetReference(line):
    import re
    reference = {}
    regex = "\(((PM|DO).*)\)"
    prog = re.compile(regex)
    result = prog.search(line)
    if result:
        ids = result.group(1)
        for id in re.split(r'[\s ]', ids):
            if id.startswith("PMID"):
                reference["PMID"] = id[5:]
            elif id.startswith("DOI"):
                doi = id[4:]
                doi = re.sub("<.*?>", "", doi)
                reference["DOI"] = doi
            elif id.startswith("PMCID"):
                reference["PMCID"] = id[6:]
        txt = line.replace("(" + ids + ")", "", -1).strip()
        if txt != "":
            reference["txt"] = txt
    else:
        txt = line.strip()
        if txt != "":
            reference["txt"] = txt
    return reference


# 补全LMSG_ID
def FullfillLMSG_ID(startID):
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")

    for result in collection_eLMSG.find({"LMSG_ID": {"$exists": 0}}, {"_id": 1}):
        startID += 1
        lmsg_id = "MSG" + str(startID).zfill(6)
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"LMSG_ID": lmsg_id}})

    CloseMongoDB(client_eLMSG)


# 测试Lineage与Name对应关系
def CheckLineage():
    import re
    ht = set()

    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")
    for result in collection_eLMSG.find({}, {"_id": 0, "Lineage": 1}):
        items = re.search("\(.*\)", result["Lineage"])
        if items != None:
            ht.add(items.group()[1:-1])
            result["Lineage"] = re.sub("\.\(.*\)", "", result["Lineage"])
        for taxon in result["Lineage"].split("."):
            if taxon.startswith("("):
                taxon = taxon[1:-1]
            ht.add(taxon)

    for taxon in ht:
        result = collection_eLMSG.find_one({"Name": taxon})
        if result == None:
            print(taxon)
    CloseMongoDB(client_eLMSG)


# 测试除Strain外的Name的唯一性
def CheckUniqueName():
    ht = set()
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")
    for result in collection_eLMSG.find({"Rank": {"$ne": "Strain"}}, {"Name": 1, "_id": 0}):
        if result["Name"] in ht:
            print(result["Name"])
        else:
            ht.add(result["Name"])
    CloseMongoDB(client_eLMSG)


# 补全LMSG_ID
def FullfillParentLMSGId():
    nameset = set(["unC", "unO", "unF", "unP", "UnC", "UnO", "UnF", "UnP"])
    ht = {}
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")
    for result in collection_eLMSG.find({"Name": {"$ne": "Strain"}}, {"Name": 1, "LMSG_ID": "1", "_id": 0}):
        ht[result["Name"]] = result["LMSG_ID"]
    for result in collection_eLMSG.find({"Parent_LMSG_ID": {"$exists": 0}}, {"Lineage": 1}):
        items = re.search("\(.*\)", result["Lineage"])
        if items != None:
            lineage = re.sub("\.\(.*\)", "", result["Lineage"]).split(".")
            for i in range(len(lineage) - 1, -1, -1):
                if lineage[i] in nameset:
                    lineage.pop(i)
            if len(lineage) > 1:
                collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"Parent_LMSG_ID": ht[lineage[-1]]}})
        else:
            lineage = result["Lineage"].split(".")
            for i in range(len(lineage) - 1, -1, -1):
                if lineage[i] in nameset:
                    lineage.pop(i)
            if len(lineage) > 2:
                collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"Parent_LMSG_ID": ht[lineage[-2]]}})
    CloseMongoDB(client_eLMSG)


# 结构化Genbank_16SrRNA
def Structure_Genbank_16SrRNA():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")
    for result in collection_eLMSG.find({"16SrRNA": {"$exists": 1}}, {"16SrRNA": 1}):
        collection_eLMSG.update_one({"_id": result["_id"]},
                                    {"$set": {"Genbank_16SrRNA": result["16SrRNA"]}, "$unset": {"16SrRNA": 1}})
    CloseMongoDB(client_eLMSG)


# 结构化JGI_IMG
def Structure_JGI_IMG():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")
    for result in collection_eLMSG.find({"JGI IMG": {"$exists": 1}}, {"JGI IMG": 1}):
        collection_eLMSG.update_one({"_id": result["_id"]},
                                    {"$set": {"JGI_IMG": result["JGI IMG"]}, "$unset": {"JGI IMG": 1}})
    CloseMongoDB(client_eLMSG)


# 结构化NCBI_Assembly
def Structure_NCBI_Assembly():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")
    for result in collection_eLMSG.find({"NCBI GenBank": {"$exists": 1}}, {"NCBI GenBank": 1}):
        collection_eLMSG.update_one({"_id": result["_id"]},
                                    {"$set": {"NCBI_Assembly": result["NCBI GenBank"]}, "$unset": {"NCBI GenBank": 1}})
    CloseMongoDB(client_eLMSG)


# 更改不标准的Type字段结构
def ChangeTypeFieldType():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")

    for result in collection_eLMSG.find({"Type": {"$exists": 1}}, {"Type": 1}):
        if isinstance(result["Type"], str):
            type = result["Type"]
            id = result["_id"]
            collection_eLMSG.update_one({"_id": id}, {"$set": {"Type": [type]}})

    CloseMongoDB(client_eLMSG)


# 更改不标准的Phenotype.Physiology.nutrition_type字段结构
def ChangeNutritionTypeFieldType():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")

    for result in collection_eLMSG.find({"Phenotype.Physiology.nutrition_type": {"$exists": 1}},
                                        {"Phenotype.Physiology.nutrition_type": 1}):
        if isinstance(result["Phenotype"]["Physiology"]["nutrition_type"], str):
            type = result["Phenotype"]["Physiology"]["nutrition_type"]
            id = result["_id"]
            collection_eLMSG.update_one({"_id": id}, {"$set": {"Phenotype.Physiology.nutrition_type": [type]}})

    CloseMongoDB(client_eLMSG)


# 更改Phenotype.Biochemistry.fattyacids字段结构
def StructureFattyAcids():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")
    for result in collection_eLMSG.find({"Phenotype.Biochemistry.fattyacids": {"$exists": 1}},
                                        {"Phenotype.Biochemistry.fattyacids": 1}):
        fattyacids = trans_fattyacids(result["Phenotype"]["Biochemistry"]["fattyacids"])
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"Phenotype.Biochemistry.fattyacids": fattyacids}})
    CloseMongoDB(client_eLMSG)


def trans_fattyacids(fattyacids_str):
    import json
    fattyacids = []
    for k, v in json.loads(fattyacids_str).items():
        fattyacids.append({"name": k, "account": str(v)})
    return fattyacids


# 增加基因组注释数据
def AddFunctionAnnotation(filepath):
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")

    infile = open(filepath, "r")
    infile.readline()
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        if splitline[0].startswith("GCA_"):
            result = collection_eLMSG.find_one({"Rank": "Strain", "NCBI_Assembly.id": splitline[0]}, {"_id": 1})
            if result != None:
                gcaId = splitline[0].replace(".", "_")
                anno = {"Swissprot": int(splitline[4]), "TIGRFAMs": int(splitline[5]), "KEGG": int(splitline[2]),
                        "COGs": int(splitline[1]),
                        "Pfam": int(splitline[3]), "GO": int(splitline[6]), "EC": int(splitline[7]),
                        "MetaCyc": int(splitline[8])}
                collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"FunctionAnnotation." + gcaId: anno}})
        else:
            result = collection_eLMSG.find_one({"Rank": "Strain", "JGI_IMG.id": splitline[0]}, {"_id": 1})
            if result != None:
                jgiId = splitline[0].replace(".", "_")
                anno = {"Swissprot": int(splitline[4]), "TIGRFAMs": int(splitline[5]), "KEGG": int(splitline[2]),
                        "COGs": int(splitline[1]),
                        "Pfam": int(splitline[3]), "GO": int(splitline[6]), "EC": int(splitline[7]),
                        "MetaCyc": int(splitline[8])}
                collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"FunctionAnnotation." + jgiId: anno}})
    infile.close()
    CloseMongoDB(client_eLMSG)


# 增加JGI_IMG元数据
def FillIMGMeta(filepath):
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")

    infile = open(filepath, "r")
    infile.readline()
    # count=2
    for line in infile:
        line = line.rstrip("\n")
        splitline = line.split("\t")
        # print(str(count)+":"+str(len(splitline)))
        # count+=1
        result = collection_eLMSG.find({"JGI_IMG.id": splitline[0].strip()[4:]}, {"_id": 1})
        if result != []:
            meta = {}
            if splitline[1].strip() != "":
                meta["JGI_IMG.$.project"] = splitline[1].strip()
            if splitline[2].strip() != "":
                meta["JGI_IMG.$.principal_investigator"] = splitline[2].strip()
            if splitline[3].strip() != "":
                meta["JGI_IMG.$.email"] = splitline[3].strip()
            if splitline[4].strip() != "":
                meta["JGI_IMG.$.sequencing_center"] = splitline[4].strip()
            if splitline[5].strip() != "":
                meta["JGI_IMG.$.submitter_date"] = splitline[5].strip()
            if splitline[6].strip() != "":
                meta["JGI_IMG.$.update_date"] = splitline[6].strip()
            if splitline[7].strip() != "":
                meta["JGI_IMG.$.submission_type"] = splitline[7].strip()
            if splitline[8].strip() != "":
                meta["JGI_IMG.$.sequencing_status"] = splitline[8].strip()
            if splitline[9].strip() != "":
                meta["JGI_IMG.$.seq_status"] = splitline[9].strip()
            if splitline[10].strip() != "":
                xref = []
                ids = splitline[10].strip().split(";")
                for id in ids:
                    if id.strip().startswith("PRJ"):
                        xref.append({"db": "bioproject", "id": id.strip()})
                    elif id.strip().startswith("SAM"):
                        xref.append({"db": "biosample", "id": id.strip()})
                if len(xref) != 0:
                    meta["JGI_IMG.$.xref"] = xref
            collection_eLMSG.update_many({"JGI_IMG": {"$elemMatch": {"id": splitline[0].strip()[4:]}}}, {"$set": meta})
    infile.close()

    CloseMongoDB(client_eLMSG)


# 创建所需索引
def CreateIndexes():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")

    index1 = IndexModel([("Name", ASCENDING)])
    index2 = IndexModel([("LMSG_ID", ASCENDING)])
    index3 = IndexModel([("Parent_LMSG_ID", ASCENDING)])
    index4 = IndexModel([("NCBI_Assembly.id", ASCENDING)])
    index5 = IndexModel([("StructuredLineage", ASCENDING)])
    index6 = IndexModel([("Rank", ASCENDING)])
    index7 = IndexModel([("Type", ASCENDING)])
    index8 = IndexModel([("JGI_IMG.id", ASCENDING)])
    index9 = IndexModel([("Phenotype.Physiology.nutrition_type", ASCENDING)])
    index10 = IndexModel([("Phenotype.Biochemistry.fattyacids", ASCENDING)])

    collection_eLMSG.create_indexes([index1, index2, index3, index4, index5, index6, index7, index8, index9, index10])

    CloseMongoDB(client_eLMSG)


def TestUniqueAssemblyID():
    ht = {}
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180612")

    results = collection_eLMSG.find({"Rank": "Strain", "NCBI_Assembly.id": {"$exists": 1}}, {"NCBI_Assembly": 1})
    for result in results:
        for item in result["NCBI_Assembly"]:
            if item["id"] in ht:
                ht[item["id"]] += 1
            else:
                ht[item["id"]] = 1
    CloseMongoDB(client_eLMSG)
    for key in ht.keys():
        if ht[key] > 1:
            print(key)


if __name__ == '__main__':
    # Clean_eLMSG()
    # CreateIndexes()
    # CheckLineage()
    # CheckUniqueName()

    # FullfillLMSG_ID(0)
    # FullfillParentLMSGId()
    # GenerateStructuredReference()
    # Structure_Genbank_16SrRNA()
    # Structure_JGI_IMG()
    # Structure_NCBI_Assembly()
    # ChangeTypeFieldType()
    # ChangeNutritionTypeFieldType()
    # StructureFattyAcids()
    FillIMGMeta("D:/IMG_genome_metadata.txt")
    # TestUniqueAssemblyID()
    AddFunctionAnnotation("D:/barplot.txt")
